#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Debian 系统 & 硬件全方位体检
运行：sudo python3 debian_hw_audit.py  [--json|--csv]
依赖：apt install python3-psutil
"""
import argparse
import csv
import datetime
import json
import os
import platform
import subprocess
import sys
from collections import OrderedDict
from pathlib import Path

# ---------- 可选依赖 ----------
try:
    import psutil
except ImportError:
    sys.exit("请先安装依赖：sudo apt update && sudo apt install python3-psutil")

# ---------- 人类可读单位 ----------
def bytes2human(n, fmt=".1f"):
    symbols = ("B", "KB", "MB", "GB", "TB")
    prefix = {}
    for i, s in enumerate(symbols):
        prefix[s] = 1 << (i * 10)
    for s in reversed(symbols[1:]):
        if n >= prefix[s]:
            return f"{n / prefix[s]:{fmt}}{s}"
    return f"{n}{symbols[0]}"

# ---------- 基础信息 ----------
def system_info():
    return OrderedDict(
        timestamp=datetime.datetime.now().isoformat(timespec="seconds"),
        debian_version=Path("/etc/debian_version").read_text().strip(),
        kernel=platform.release(),
        arch=platform.machine(),
        hostname=platform.node(),
        uptime=subprocess.check_output(["uptime", "-p"], text=True).strip(),
    )

# ---------- CPU ----------
def cpu_info():
    cpu_percent = psutil.cpu_percent(interval=1)
    logical = psutil.cpu_count()
    physical = psutil.cpu_count(logical=False)
    freq = psutil.cpu_freq()
    return OrderedDict(
        cpu_physical_cores=physical,
        cpu_logical_cores=logical,
        cpu_current_freq_MHz=round(freq.current),
        cpu_max_freq_MHz=round(freq.max),
        cpu_load_percent=round(cpu_percent, 1),
        cpu_load_health="OK" if cpu_percent < 80 else "HIGH",
    )

# ---------- 内存 ----------
def mem_info():
    vm = psutil.virtual_memory()
    swap = psutil.swap_memory()
    return OrderedDict(
        mem_total=bytes2human(vm.total),
        mem_available=bytes2human(vm.available),
        mem_used_percent=vm.percent,
        mem_health="OK" if vm.percent < 80 else "HIGH",
        swap_total=bytes2human(swap.total),
        swap_used_percent=swap.percent,
    )

# ---------- 磁盘 ----------
def disk_info():
    parts = []
    for p in psutil.disk_partitions():
        if p.fstype in ("squashfs", "tmpfs"):
            continue
        usage = psutil.disk_usage(p.mountpoint)
        parts.append(
            OrderedDict(
                device=p.device,
                mountpoint=p.mountpoint,
                fstype=p.fstype,
                size_total=bytes2human(usage.total),
                size_used=bytes2human(usage.used),
                used_percent=round(usage.used / usage.total * 100, 1),
                health="OK" if usage.used / usage.total < 0.9 else "HIGH",
            )
        )
    return parts

# ---------- 温度 & 风扇（如有） ----------
def thermal_info():
    try:
        temps = psutil.sensors_temperatures()
    except AttributeError:
        temps = {}
    try:
        fans = psutil.sensors_fans()
    except AttributeError:
        fans = {}
    data = OrderedDict()
    for name, entries in temps.items():
        data[f"temp_{name}_C"] = round(entries[0].current, 1)
    for name, entries in fans.items():
        data[f"fan_{name}_RPM"] = entries[0].current
    return data or {"thermal": "no sensor"}

# ---------- 电池（笔记本/灵芯派） ----------
def battery_info():
    if not psutil.sensors_battery():
        return {"battery_present": False}
    batt = psensors_battery()
    return OrderedDict(
        battery_percent=round(batt.percent, 1),
        battery_plugged=batt.power_plugged,
        battery_secs_left=batt.secsleft if batt.secsleft != psutil.POWER_TIME_UNLIMITED else "unlimited",
    )

# ---------- 灵芯派专用 ----------
def smartpi_info():
    try:
        from smartpi import base_driver
    except ImportError:
        return {"smartpi": "library not found"}
    base_driver.smartpi_init()
    return OrderedDict(
        device_model=base_driver.read_device_model(),
        firmware_version=base_driver.read_version(),
        factory_data=base_driver.read_factory_data(),
        hardware_id=base_driver.read_hardware_ID(),
        device_name=base_driver.read_device_name(),
        connected=base_driver.read_connected(),
        battery=base_driver.read_battery(),
    )


# ---------- 汇总 ----------
def collect_all():
    data = OrderedDict()
    data.update(system_info())
    data.update(cpu_info())
    data.update(mem_info())
    data["disks"] = disk_info()
    data.update(thermal_info())
    data.update(battery_info())
    data.update(smartpi_info())
    return data


# ---------- 输出 ----------
def main():
    parser = argparse.ArgumentParser(description="Debian 硬件/系统体检")
    parser.add_argument("--json", action="store_true", help="导出 JSON")
    parser.add_argument("--csv", action="store_true", help="导出 CSV")
    args = parser.parse_args()

    data = collect_all()

    if args.json:
        print(json.dumps(data, indent=2, ensure_ascii=False))
        return
    if args.csv:
        # 扁平化一级，disks 单独展开
        flat = data.copy()
        disks = flat.pop("disks")
        writer = csv.DictWriter(sys.stdout, fieldnames=flat.keys())
        writer.writeheader()
        writer.writerow(flat)
        if disks:
            writer = csv.DictWriter(sys.stdout, fieldnames=disks[0].keys())
            writer.writeheader()
            writer.writerows(disks)
        return

    # 默认：人类可读
    def print_kv(d, indent=0):
        prefix = " " * indent
        for k, v in d.items():
            if isinstance(v, list):
                for idx, item in enumerate(v, 1):
                    print(f"{prefix}{k} #{idx}")
                    print_kv(item, indent + 2)
            else:
                print(f"{prefix}{k:<25}: {v}")

    print("========== Debian & 硬件全方位体检 ==========")
    print_kv(data)
    print("============================================")


if __name__ == "__main__":
    main()